#! /usr/bin/env python
#
# Implement cookbook using pyfitsio
import sys
import os

import ctypes

# Use just the C-types bindings to CFITSIO
from fitsio import cfitsio, pyfitsio

def put_table_data(ff, planet, diameter, density):

    nrows = len(planet)

    status = ctypes.c_int(0)

    tfields = ctypes.c_int(3)
    trows   = ctypes.c_longlong(nrows)

    firstrow  = ctypes.c_longlong(1)
    firstelem = ctypes.c_longlong(1)

    cfitsio.fits_write_col(ff.fptr, cfitsio.TSTRING, 1, 
                           firstrow, firstelem, nrows,
                           (ctypes.c_char_p * nrows)(*planet), 
                           ctypes.byref(status))

    cfitsio.fits_write_col(ff.fptr, cfitsio.TLONG, 2, 
                           firstrow, firstelem, nrows,
                           (ctypes.c_long * nrows)(*diameter), 
                           ctypes.byref(status))

    cfitsio.fits_write_col(ff.fptr, cfitsio.TFLOAT, 3, 
                           firstrow, firstelem, nrows,
                           (ctypes.c_float * nrows)(*density), 
                           ctypes.byref(status))

    pyfitsio.check_error(status)

def main():
    '''
    Create an ASCII table extension containing 3 columns and 6 rows
    '''

    # Put in some data
    planet = ['Mercury', 'Venus', 'Earth', 'Mars', 'Jupiter', 'Saturn']
    diameter = [4880, 12112, 12742, 6800, 143000, 121000]
    density = [5.1, 5.3, 5.52, 3.94, 1.33, 0.69]

    filename = 'atestfil.fit'
    extname = 'PLANETS_ASCII'

    # Column names, formats, and units
    ttype = ['Planet', 'Diameter', 'Density']
    tform = ['a8', 'I6', 'F4.2']
    tunit = ['\0', 'km', 'g/cm']
    nrows = 6

    # Open file for reading and writing
    ff = pyfitsio.open(filename, mode='rw')

    ff.create_table(cfitsio.ASCII_TBL, extname, ttype, tform, tunit, nrows)

    put_table_data(ff, planet, diameter, density)

    ff.close()

if __name__ == '__main__':
    main()
